#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库连接管理模块

提供 MySQL 数据库连接池管理、查询执行和事务处理功能。
支持自定义端口配置和完善的错误处理机制。
"""

import time
from typing import Optional, Dict, Any, List, Tuple, Union
from contextlib import contextmanager
import mysql.connector
from mysql.connector import pooling, Error as MySQLError
from mysql.connector.cursor import MySQLCursor
from mysql.connector.abstracts import MySQLConnectionAbstract

from .config import DatabaseConfig
from .logger import (
    get_logger,
    log_database_connection,
    log_database_connection_success,
    log_database_connection_error,
    log_query_execution,
    log_query_result,
    log_query_error,
)


class DatabaseManager:
    """
    数据库管理器
    
    负责管理 MySQL 数据库连接池，提供查询执行、事务处理等功能。
    支持自定义端口设置和完善的错误处理。
    """
    
    def __init__(self, config: DatabaseConfig):
        """
        初始化数据库管理器
        
        Args:
            config: 数据库配置对象
        """
        self.config = config
        self.logger = get_logger("database")
        self._pool: Optional[pooling.MySQLConnectionPool] = None
        self._is_connected = False
    
    def connect(self) -> None:
        """
        建立数据库连接池
        
        Raises:
            MySQLError: 当数据库连接失败时
        """
        if self._is_connected:
            self.logger.warning("数据库已连接，跳过重复连接")
            return
        
        start_time = time.time()
        
        try:
            # 记录连接信息
            log_database_connection(
                host=self.config.host,
                port=self.config.port,
                database=self.config.database,
                user=self.config.user
            )
            
            # 构建连接配置
            pool_config = {
                'pool_name': 'mysql_mcp_pool',
                'pool_size': self.config.pool_size,
                'pool_reset_session': True,
                'host': self.config.host,
                'port': self.config.port,
                'user': self.config.user,
                'password': self.config.password,
                'database': self.config.database,
                'charset': 'utf8mb4',
                'collation': 'utf8mb4_unicode_ci',
                'autocommit': True,
                'connect_timeout': self.config.connection_timeout,
            }
            
            # SSL 配置
            if not self.config.ssl_disabled:
                pool_config['ssl_disabled'] = False
                pool_config['ssl_verify_cert'] = False
                pool_config['ssl_verify_identity'] = False
            else:
                pool_config['ssl_disabled'] = True
            
            # 创建连接池
            self._pool = pooling.MySQLConnectionPool(**pool_config)
            
            # 测试连接
            self._test_connection()
            
            self._is_connected = True
            connection_time = time.time() - start_time
            log_database_connection_success(connection_time)
            
        except MySQLError as e:
            log_database_connection_error(e)
            raise
        except Exception as e:
            log_database_connection_error(e)
            raise MySQLError(f"数据库连接失败: {str(e)}") from e
    
    def _test_connection(self) -> None:
        """
        测试数据库连接
        
        Raises:
            MySQLError: 当连接测试失败时
        """
        if not self._pool:
            raise MySQLError("连接池未初始化")
        
        try:
            with self._get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT 1")
                result = cursor.fetchone()
                cursor.close()
                
                if result != (1,):
                    raise MySQLError("连接测试失败：返回结果异常")
                    
        except MySQLError:
            raise
        except Exception as e:
            raise MySQLError(f"连接测试失败: {str(e)}") from e
    
    @contextmanager
    def _get_connection(self):
        """
        获取数据库连接的上下文管理器
        
        Yields:
            MySQLConnectionAbstract: 数据库连接对象
            
        Raises:
            MySQLError: 当获取连接失败时
        """
        if not self._is_connected or not self._pool:
            raise MySQLError("数据库未连接")
        
        connection = None
        try:
            connection = self._pool.get_connection()
            yield connection
        except MySQLError:
            raise
        except Exception as e:
            raise MySQLError(f"获取数据库连接失败: {str(e)}") from e
        finally:
            if connection and connection.is_connected():
                connection.close()
    
    def execute_query(
        self,
        query: str,
        params: Optional[Union[Dict[str, Any], Tuple, List]] = None,
        fetch_all: bool = True
    ) -> List[Dict[str, Any]]:
        """
        执行 SELECT 查询
        
        Args:
            query: SQL 查询语句
            params: 查询参数
            fetch_all: 是否获取所有结果
            
        Returns:
            List[Dict[str, Any]]: 查询结果列表
            
        Raises:
            MySQLError: 当查询执行失败时
        """
        start_time = time.time()
        
        try:
            log_query_execution(query, params)
            
            with self._get_connection() as conn:
                cursor = conn.cursor(dictionary=True)
                
                if params:
                    cursor.execute(query, params)
                else:
                    cursor.execute(query)
                
                if fetch_all:
                    results = cursor.fetchall()
                else:
                    results = [cursor.fetchone()] if cursor.fetchone() else []
                
                row_count = len(results) if results else 0
                cursor.close()
                
                execution_time = time.time() - start_time
                log_query_result(query, row_count, execution_time)
                
                return results or []
                
        except MySQLError as e:
            log_query_error(query, e)
            raise
        except Exception as e:
            error = MySQLError(f"查询执行失败: {str(e)}")
            log_query_error(query, error)
            raise error from e
    
    def execute_update(
        self,
        query: str,
        params: Optional[Union[Dict[str, Any], Tuple, List]] = None
    ) -> int:
        """
        执行 INSERT/UPDATE/DELETE 语句
        
        Args:
            query: SQL 语句
            params: 查询参数
            
        Returns:
            int: 受影响的行数
            
        Raises:
            MySQLError: 当执行失败时
        """
        start_time = time.time()
        
        try:
            log_query_execution(query, params)
            
            with self._get_connection() as conn:
                cursor = conn.cursor()
                
                if params:
                    cursor.execute(query, params)
                else:
                    cursor.execute(query)
                
                affected_rows = cursor.rowcount
                conn.commit()
                cursor.close()
                
                execution_time = time.time() - start_time
                log_query_result(query, affected_rows, execution_time)
                
                return affected_rows
                
        except MySQLError as e:
            log_query_error(query, e)
            raise
        except Exception as e:
            error = MySQLError(f"更新执行失败: {str(e)}")
            log_query_error(query, error)
            raise error from e
    
    def execute_many(
        self,
        query: str,
        params_list: List[Union[Dict[str, Any], Tuple, List]]
    ) -> int:
        """
        批量执行 SQL 语句
        
        Args:
            query: SQL 语句
            params_list: 参数列表
            
        Returns:
            int: 受影响的总行数
            
        Raises:
            MySQLError: 当执行失败时
        """
        start_time = time.time()
        
        try:
            self.logger.debug(
                "批量执行 SQL",
                query=query,
                batch_size=len(params_list)
            )
            
            with self._get_connection() as conn:
                cursor = conn.cursor()
                cursor.executemany(query, params_list)
                
                affected_rows = cursor.rowcount
                conn.commit()
                cursor.close()
                
                execution_time = time.time() - start_time
                self.logger.info(
                    "批量执行完成",
                    query_preview=query[:100] + "..." if len(query) > 100 else query,
                    batch_size=len(params_list),
                    affected_rows=affected_rows,
                    execution_time=f"{execution_time:.3f}s"
                )
                
                return affected_rows
                
        except MySQLError as e:
            log_query_error(query, e)
            raise
        except Exception as e:
            error = MySQLError(f"批量执行失败: {str(e)}")
            log_query_error(query, error)
            raise error from e
    
    @contextmanager
    def transaction(self):
        """
        事务上下文管理器
        
        Yields:
            MySQLConnectionAbstract: 数据库连接对象
            
        Raises:
            MySQLError: 当事务执行失败时
        """
        if not self._is_connected or not self._pool:
            raise MySQLError("数据库未连接")
        
        connection = None
        try:
            connection = self._pool.get_connection()
            connection.autocommit = False
            connection.start_transaction()
            
            self.logger.debug("开始事务")
            yield connection
            
            connection.commit()
            self.logger.debug("事务提交成功")
            
        except MySQLError as e:
            if connection:
                connection.rollback()
                self.logger.error("事务回滚", error=str(e))
            raise
        except Exception as e:
            if connection:
                connection.rollback()
                self.logger.error("事务回滚", error=str(e))
            raise MySQLError(f"事务执行失败: {str(e)}") from e
        finally:
            if connection and connection.is_connected():
                connection.autocommit = True
                connection.close()
    
    def get_table_info(self, table_name: str) -> Dict[str, Any]:
        """
        获取表结构信息
        
        Args:
            table_name: 表名
            
        Returns:
            Dict[str, Any]: 表结构信息
            
        Raises:
            MySQLError: 当获取表信息失败时
        """
        try:
            # 获取表结构
            columns_query = """
                SELECT 
                    COLUMN_NAME,
                    DATA_TYPE,
                    IS_NULLABLE,
                    COLUMN_DEFAULT,
                    COLUMN_KEY,
                    EXTRA,
                    COLUMN_COMMENT
                FROM INFORMATION_SCHEMA.COLUMNS 
                WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
                ORDER BY ORDINAL_POSITION
            """
            
            columns = self.execute_query(
                columns_query,
                (self.config.database, table_name)
            )
            
            # 获取索引信息
            indexes_query = """
                SELECT 
                    INDEX_NAME,
                    COLUMN_NAME,
                    NON_UNIQUE,
                    INDEX_TYPE
                FROM INFORMATION_SCHEMA.STATISTICS 
                WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
                ORDER BY INDEX_NAME, SEQ_IN_INDEX
            """
            
            indexes = self.execute_query(
                indexes_query,
                (self.config.database, table_name)
            )
            
            return {
                'table_name': table_name,
                'columns': columns,
                'indexes': indexes
            }
            
        except MySQLError:
            raise
        except Exception as e:
            raise MySQLError(f"获取表信息失败: {str(e)}") from e
    
    def get_database_info(self) -> Dict[str, Any]:
        """
        获取数据库信息
        
        Returns:
            Dict[str, Any]: 数据库信息
            
        Raises:
            MySQLError: 当获取数据库信息失败时
        """
        try:
            # 获取数据库版本
            version_result = self.execute_query("SELECT VERSION() as version")
            version = version_result[0]['version'] if version_result else 'Unknown'
            
            # 获取表列表
            tables_query = """
                SELECT 
                    TABLE_NAME,
                    TABLE_TYPE,
                    TABLE_COMMENT
                FROM INFORMATION_SCHEMA.TABLES 
                WHERE TABLE_SCHEMA = %s
                ORDER BY TABLE_NAME
            """
            
            tables = self.execute_query(tables_query, (self.config.database,))
            
            return {
                'database_name': self.config.database,
                'version': version,
                'host': self.config.host,
                'port': self.config.port,
                'tables': tables
            }
            
        except MySQLError:
            raise
        except Exception as e:
            raise MySQLError(f"获取数据库信息失败: {str(e)}") from e
    
    def disconnect(self) -> None:
        """
        断开数据库连接
        """
        if self._is_connected and self._pool:
            try:
                # 关闭连接池中的所有连接
                # 注意：mysql-connector-python 的连接池没有直接的关闭方法
                # 我们只需要将引用设为 None，让垃圾回收器处理
                self._pool = None
                self._is_connected = False
                self.logger.info("数据库连接已断开")
            except Exception as e:
                self.logger.error("断开数据库连接时发生错误", error=str(e))
    
    @property
    def is_connected(self) -> bool:
        """
        检查是否已连接到数据库
        
        Returns:
            bool: 连接状态
        """
        return self._is_connected and self._pool is not None
    
    def __enter__(self):
        """
        上下文管理器入口
        
        Returns:
            DatabaseManager: 数据库管理器实例
        """
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        上下文管理器出口
        
        Args:
            exc_type: 异常类型
            exc_val: 异常值
            exc_tb: 异常追踪
        """
        self.disconnect()